#include <assert.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <cusolverDn.h>

#include "sdm.h"
#include <Utility/cuda_device.cuh>

namespace shermit {
namespace fitting {
namespace curve {
	template<typename Float_t>
	__global__ void calcSdmMetricsMatrixClosed(
		CubicBSplineCurve<Float_t, DeviceType::GPU> curve,
		int nd, int nc0,
		const Float_t* t,
		const nv::vec2<Float_t> *p_tk,
		const nv::vec2<Float_t> *dpts,
		const nv::vec2<Float_t> *tangents,
		const nv::vec2<Float_t>* normals,
		double* matrix) 
	{
		int idx = blockIdx.x * blockDim.x + threadIdx.x;
		if (idx >= nd * nc0) return;
		int i = idx / nc0;      // column index
		int j = idx - i * nc0;  // row index
		Float_t base = curve.baseFuncVal(j, t[i]);
		if (j < 3) {
			base += curve.baseFuncVal(nc0 + j, t[i]);
		}
		// calculate error term d/(d-rho) * Tk*Tk^t + Nk*Nk^t if d<0 else Nk*Nk^t
		
		auto v = dpts[i] - p_tk[i];
		const auto nk = normals[i];
		Float_t w0 = nk.x * nk.x;
		Float_t w1 = nk.x * nk.y;
		Float_t w3 = nk.y * nk.y;
		Float_t side = nv::dot(v, nk);
		Float_t d = nv::length(v);
		Float_t rho = 1.0 / curve.curvature(t[i]);
		
		if (side >0 ) {
			Float_t coef = (d) / ((d) + rho);
			const auto tk = tangents[i];
			w0 += coef * tk.x * tk.x;
			w1 += coef * tk.x * tk.y;
			w3 += coef * tk.y * tk.y;
		}

		int mid = 4 * nc0 * i + 2 * j;
		matrix[mid] =  base*w0;
		matrix[mid + 1] = base*w1;
		matrix[mid + 2 * nc0] = w1;
		matrix[mid + 2 * nc0 + 1] =  base*w3;
	}

	template<typename Float_t>
	__global__ void calcTdmMetricsMatrixClosed(
		CubicBSplineCurve<Float_t, DeviceType::GPU> curve,
		int nd, int nc0,
		const Float_t* t,
		const nv::vec2<Float_t>* p_tk,
		const nv::vec2<Float_t>* dpts,
		const nv::vec2<Float_t>* tangents,
		const nv::vec2<Float_t>* normals,
		double* matrix)
	{
		int idx = blockIdx.x * blockDim.x + threadIdx.x;
		if (idx >= nd * nc0) return;
		int i = idx / nc0;      // column index
		int j = idx - i * nc0;  // row index
		Float_t base = curve.baseFuncVal(j, t[i]);
		if (j < 3) {
			base += curve.baseFuncVal(nc0 + j, t[i]);
		}
		// calculate error term d/(d-rho) * Tk*Tk^t + Nk*Nk^t if d<0 else Nk*Nk^t

		auto v = dpts[i] - p_tk[i];
		const auto nk = normals[i];
		Float_t w0 = nk.x * nk.x;
		Float_t w1 = nk.x * nk.y;
		Float_t w3 = nk.y * nk.y;

		int mid = 4 * nc0 * i + 2 * j;
		matrix[mid] = base * w0;
		matrix[mid + 1] = w1;
		matrix[mid + 2 * nc0] = w1;
		matrix[mid + 2 * nc0 + 1] = base * w3;
	}

	template<typename Float_t>
	__global__ void calcBfmExpandedClosed(
		CubicBSplineCurve<Float_t, DeviceType::GPU> curve,
		int nd, int nc0,
		const double* t,
		double* matrix
	)
	{
		int idx = blockIdx.x * blockDim.x + threadIdx.x;
		if (idx >= nd * nc0) return;
		int i = idx / nd;      // column index
		int j = idx - i * nd;  // row index
		int mid = 4 * nd * i + 2 * j;
		Float_t base = curve.baseFuncVal(i, t[j]);
		if (i < 3) {
			base += curve.baseFuncVal(nc0 + i, t[j]);
		}
		matrix[mid] = base;
		matrix[mid + 2 * nd + 1] = base;
	}

	void CudaSDMSolverLU64::solve(double error_target, DistanceMetrics metrics) {
		const int m = 2 * nd;
		const int nc_folded = nc - 3;
		const int n = 2 * nc_folded;
		const int lda0 = n;
		const int lda1 = m;
		const int lda0a1 = n;
		const int ldb = m;
		ipiv.resize(n);
		A0.resize(n * m);
		A1.resize(m * n);
		A0A1.resize(n * n);
		b.resize(m);
		thrust::device_vector<int> info(1, 0);

		double* buffer = getRawPtr(workspace);
		int* ipiv_ptr = getRawPtr(ipiv);
		int* info_ptr = getRawPtr(info);
		double* A0_ptr = getRawPtr(A0);
		double* A1_ptr = getRawPtr(A1);
		double* A0A1_ptr = getRawPtr(A0A1);
		double* b_ptr = getRawPtr(b);

		nv::vec2<double>* pts_ptr= const_cast<nv::vec2<double>*>(projector.getDataPoints());
		double* dpts_ptr = (double*)pts_ptr;
		auto& CudaDevice = MCudaDevice::getInstance();

		cublasHandle_t cublas_h = NULL;
		cusolverDnHandle_t cusolver_h = NULL;
		cusolverDnCreate(&cusolver_h);
		cublasCreate(&cublas_h);

		int _iter = 0;
		double rmse;
		double* t_ptr;
		nv::vec2<double>* pts_projected, * tangents, * normals;
		while (_iter < max_iter_solver) {
			projector.computeProjection(curve, max_iter_projection, spr, &rmse, &t_ptr, &pts_projected, &tangents, &normals);
			if (rmse < error_target) break;

			A0A1.assign(n * n, 0);
			A1.assign(m * n, 0);
			b.assign(m, 0);

			std::pair<int, int> gridblk;
			if (metrics == DistanceMetrics::SDM) {
				gridblk = CudaDevice.getOptLaunchParams(nd * nc_folded, 0, calcSdmMetricsMatrixClosed<double>);
				calcSdmMetricsMatrixClosed << <gridblk.first, gridblk.second >> >
					(curve, nd, nc_folded, t_ptr, pts_projected, pts_ptr, tangents, normals, getRawPtr(A0));
			}
			else if (metrics == DistanceMetrics::TDM) {
				gridblk = CudaDevice.getOptLaunchParams(nd * nc_folded, 0, calcTdmMetricsMatrixClosed<double>);
				calcTdmMetricsMatrixClosed << <gridblk.first, gridblk.second >> >
					(curve, nd, nc_folded, t_ptr, pts_projected, pts_ptr, tangents, normals, getRawPtr(A0));
			}
			else {
				printf("unsupported metrics!\n");
				abort();
			}

			gridblk = CudaDevice.getOptLaunchParams(nd* nc_folded, 0, calcBfmExpandedClosed<double>);
			calcBfmExpandedClosed << <gridblk.first, gridblk.second >> >
				(curve, nd, nc_folded, t_ptr, getRawPtr(A1));
			cudaDeviceSynchronize();
			getLastCudaError("failure when computing A0 or A1");


			double alpha = 1;
			double beta = 1;
			cublasStatus_t status;
			// A0A1 = A0*A1
			status = cublasDgemm(cublas_h, 
				CUBLAS_OP_N, CUBLAS_OP_N, n, n, m, 
				&alpha, A0_ptr, lda0, A1_ptr, lda1, &beta, 
				A0A1_ptr, lda0a1);

			if (status != CUBLAS_STATUS_SUCCESS) {
				printf("Matrix multiply A0*A1 failed.\n");
				cusolverDnDestroy(cusolver_h);
				cublasDestroy(cublas_h);
				abort();
			}
			// prepare right side of the equation A0*b
			status = cublasDgemv(cublas_h, CUBLAS_OP_N, n, m, &alpha, A0_ptr, lda0, dpts_ptr, 1, &beta, b_ptr, 1);

			if (status != CUBLAS_STATUS_SUCCESS) {
				printf("Matrix multiply A0*b failed.\n");
				cusolverDnDestroy(cusolver_h);
				cublasDestroy(cublas_h);
				abort();
			}

			int buffer_size;
			cusolverDnDgetrf_bufferSize(cusolver_h, n, n, A0A1_ptr, lda0a1, &buffer_size);
			workspace.resize(buffer_size);

			cusolverDnDgetrf(cusolver_h, n, n, A0A1_ptr, lda0a1, buffer, ipiv_ptr, info_ptr);
			int h_info = 1;
			checkCudaErrors(GPU2CPU(&h_info, info_ptr, sizeof(int)));
			if (h_info != 0) {
				printf("Error: LU factorization failed.\n");
				if (h_info < 0) printf("%dth parameter is wrong.\n", -h_info);
				else printf("U[%d, %d]=0\n", h_info, h_info);
				cusolverDnDestroy(cusolver_h);
				cublasDestroy(cublas_h);
				abort();
			}
			auto solver_status = cusolverDnDgetrs(
				cusolver_h, CUBLAS_OP_N, n, 1,
				A0A1_ptr, lda0a1, ipiv_ptr,
				b_ptr, ldb, info_ptr);
			if (solver_status != CUSOLVER_STATUS_SUCCESS) {
				printf("Failed to solve A0A1*x=A0*b.\n");
				cusolverDnDestroy(cusolver_h);
				cublasDestroy(cublas_h);
				abort();
			}
			checkCudaErrors(GPU2CPU(&h_info, info_ptr, sizeof(int)));
			if (h_info != 0) {
				printf("Error: solve x = A0A1 \ (A0*b) failed.\n");
				if (h_info < 0) printf("%dth parameter is wrong.\n", -h_info);
				else printf("U[%d, %d]=0\n", h_info, h_info);
				cusolverDnDestroy(cusolver_h);
				cublasDestroy(cublas_h);
				abort();
			}
			checkCudaErrors(GPU2GPU(curve.control_pts, b_ptr, n * sizeof(double)));
			checkCudaErrors(GPU2GPU(curve.control_pts + nc_folded, b_ptr, 6 * sizeof(double)));
			_iter++;
		}
		printf("SDM Solver iter: %d, RMSE: %lf\n", _iter, rmse);
		cusolverDnDestroy(cusolver_h);
		cublasDestroy(cublas_h);
		return;
	}
}
}
}